import os
import h5py
import torch
from torch.utils.data import Dataset, DataLoader

class AVEDataset(Dataset):
    """
    AVE数据集的自定义Dataset类，用于加载视听特征数据及对应的标签
    继承自PyTorch的Dataset类，实现了自定义数据加载逻辑
    """
    def __init__(self, data_root, split='train'):
        """
        初始化数据集，设置数据路径并区分训练/测试模式
        
        参数:
            data_root (str): 数据根目录，包含所有特征和标签文件
            split (str): 数据集划分，可选'train'（训练集）或'test'（测试集），默认为'train'
        """
        super(AVEDataset, self).__init__()
        self.split = split  # 保存数据集划分类型

        # 干净样本（含目标信号）的特征路径
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')

        # 噪声样本（仅含背景，无目标）的特征路径
        self.noisy_visual_feature_path = os.path.join(data_root, 'visual_feature_noisy.h5')  # 仅背景的视觉特征
        self.noisy_audio_feature_path = os.path.join(data_root, 'audio_feature_noisy.h5')    # 仅背景的音频特征

        # 标签路径
        self.labels_path = os.path.join(data_root, 'labels.h5')              # 原始标签（用于测试）
        self.dir_labels_path = os.path.join(data_root, 'mil_labels.h5')       # 视频级标签（视频级标签）
        self.dir_labels_bg_path = os.path.join(data_root, 'labels_noisy.h5')  # 噪声样本标签（仅背景）

        # 当前划分的样本顺序文件路径（控制训练/测试样本的选取）
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')

        self.h5_isOpen = False  # 标记h5文件是否已打开（延迟加载优化）

    def __getitem__(self, index):
        """
        根据索引获取单个样本（支持干净样本和噪声样本）
        参数:
            index (int): 样本索引（训练时包含干净样本+噪声样本的总索引）
            
        返回:
            visual_feat: 视觉特征（干净或噪声）
            audio_feat: 音频特征（干净或噪声）
            label: 对应标签（干净样本标签或噪声样本标签）
        """
        # 延迟打开h5文件：首次访问时才加载，节省资源
        if not self.h5_isOpen:
            # 加载样本顺序（用于定位干净样本）
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 加载干净样本的视听特征和标签
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            self.clean_labels = h5py.File(self.dir_labels_path, 'r')['avadataset']

            # 训练集需要加载噪声样本（负样本）
            if self.split == 'train':
                self.negative_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']
                self.negative_visual_feature = h5py.File(self.noisy_visual_feature_path, 'r')['avadataset']
                self.negative_audio_feature = h5py.File(self.noisy_audio_feature_path, 'r')['avadataset']

            # 测试集需要加载原始标签（用于评估）
            if self.split == 'test':
                self.labels = h5py.File(self.labels_path, 'r')['avadataset']

            # 标记文件已打开
            self.h5_isOpen = True  

        # 计算干净样本的数量（用于区分索引对应的是干净样本还是噪声样本）
        clean_length = len(self.sample_order)

        # 索引超过干净样本数量时，取噪声样本（仅训练时有效）
        if index >= clean_length:
            valid_index = index - clean_length  # 转换为噪声样本的本地索引
            visual_feat = self.negative_visual_feature[valid_index]
            audio_feat = self.negative_audio_feature[valid_index]
            label = self.negative_labels[valid_index]
        else:
            # 索引在干净样本范围内：训练时取视频级标签，测试时取原始标签
            sample_index = self.sample_order[index]  # 获取干净样本的实际索引
            visual_feat = self.visual_feature[sample_index]
            audio_feat = self.audio_feature[sample_index]
            if self.split == 'train':
                label = self.clean_labels[sample_index]  # 训练用视频级标签
            else:
                label = self.labels[sample_index]  # 测试用原始标签

        return visual_feat, audio_feat, label

    def __len__(self):
        """
        获取当前划分的样本总数（训练集包含干净样本+噪声样本，测试集仅含干净样本）
        
        返回:
            length (int): 样本总数
        """
        if self.split == 'train':
            # 训练集长度 = 干净样本数 + 噪声样本数
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            noisy_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']
            length = len(sample_order) + len(noisy_labels)
        elif self.split == 'test':
            # 测试集长度 = 干净样本数
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            length = len(sample_order)
        else:
            # 不支持其他划分
            raise NotImplementedError(f"不支持的数据集划分: {self.split}")

        return length

